package nsfs

import (
	"GLibHac/crypto"
	"GLibHac/fssystem"
	"GLibHac/tools"
	"encoding/binary"
	"fmt"
	"github.com/pkg/errors"
)

type NcaSectionEntry struct {
	StartBlock int32
	EndBlock   int32
}

func NewNcaSectionEntry(data []byte) *NcaSectionEntry {
	return &NcaSectionEntry{
		StartBlock: tools.BytesToIntLittleEndian(data),
		EndBlock:   tools.BytesToIntLittleEndian(data[4:]),
	}
}

const (
	RightsIdOffset       = 0x230
	RightsIdSize         = 0x10
	SectionEntriesOffset = 0x240
	SectionEntrySize     = 0x10
	FsHeaderHashOffset   = 0x280
	FsHeaderHashSize     = 0x20
	KeyAreaOffset        = 0x300
	KeyAreaSize          = 0x100
	FsHeadersOffset      = 0x400
	FsHeaderSize         = 0x200
	SectionCount         = 4
)

var Nca0FixedBodyKeySha256Hash = []byte{
	0x9A, 0xBB, 0xD2, 0x11, 0x86, 0x00, 0x21, 0x9D, 0x7A, 0xDC, 0x5B, 0x43, 0x95, 0xF8, 0x4E, 0xFD,
	0xFF, 0x6B, 0x25, 0xEF, 0x9F, 0x96, 0x85, 0x28, 0x18, 0x9E, 0x76, 0xB0, 0x92, 0xF0, 0x6A, 0xCB,
}

var NcaHeaderKeyIncorrectErr = errors.New("unable to decrypt NCA header. The file is not an NCA file or the header key is incorrect")
var InvalidFsHeaderHashErr = errors.New("FS header hash is invalid")

const NcaHeaderSize = 0xC00
const NcaHeaderSectorSize = 0x200

type DistributionType byte

const (
	DistributionTypeDownload = DistributionType(iota)
	DistributionTypeGameCard
)

type NcaContentType byte

func (n NcaContentType) String() string {
	switch n {
	case NcaContentTypeProgram:
		return "Program"
	case NcaContentTypeMeta:
		return "Meta"
	case NcaContentTypeControl:
		return "Control"
	case NcaContentTypeManual:
		return "Manual"
	case NcaContentTypeData:
		return "Data"
	case NcaContentTypePublicData:
		return "PublicData"
	default:
		return fmt.Sprintf("NcaContentType(%d)", n)
	}
}

const (
	NcaContentTypeProgram = NcaContentType(iota)
	NcaContentTypeMeta
	NcaContentTypeControl
	NcaContentTypeManual
	NcaContentTypeData
	NcaContentTypePublicData
)

type NcaVersion byte

func (n NcaVersion) String() string {
	switch n {
	case NcaVersionUnknown:
		return "Unknown"
	case NcaVersionNca3:
		return "NCA3"
	case NcaVersionNca2:
		return "NCA2"
	case NcaVersionNca0:
		return "NCA0"
	case NcaVersionNca0FixedKey:
		return "NCA0FixedKey"
	case NcaVersionNca0RsaOaep:
		return "NCA0RsaOaep"
	default:
		return fmt.Sprintf("NcaVersion(%d)", n)
	}
}

const (
	NcaVersionUnknown = NcaVersion(iota)
	NcaVersionNca3
	NcaVersionNca2
	NcaVersionNca0
	NcaVersionNca0FixedKey
	NcaVersionNca0RsaOaep
)

type NcaHeader struct {
	IsEncrypted   bool
	FormatVersion NcaVersion

	Signature1       []byte
	Signature2       []byte
	Magic            string
	DistributionType DistributionType
	ContentType      NcaContentType
	KeyGeneration1   byte
	KeyAreaKeyIndex  byte
	NcaSize          uint64
	TitleId          uint64
	ContentIndex     int32
	SdkVersion       uint32
	KeyGeneration2   byte
	RightsId         []byte
	HeaderData       []byte
}

func (h *NcaHeader) IsSectionEnabled(index int32) bool {
	info := h.GetSectionEntry(index)
	sectStart := info.StartBlock
	sectSize := info.EndBlock - sectStart
	return sectStart != 0 || sectSize != 0
}

func (h *NcaHeader) GetSectionEntry(index int32) *NcaSectionEntry {
	return NewNcaSectionEntry(h.HeaderData[SectionEntriesOffset+SectionEntrySize*index:])
}

func (h *NcaHeader) GetFsHeaderHash(index int32) []byte {
	offset := FsHeaderHashOffset + FsHeaderHashSize*index
	return h.HeaderData[offset : offset+FsHeaderHashSize]
}

func (h *NcaHeader) GetSectionStartOffset(index int32) int64 {
	return h.BlockToOffset(h.GetSectionEntry(index).StartBlock)
}

func (h *NcaHeader) HasRightsId() bool {
	return !tools.IsZeroArray(h.RightsId)
}

func (h *NcaHeader) GetEncryptedKey(index int) ([]byte, error) {
	if index < 0 || index >= SectionCount {
		return nil, errors.Errorf("key index must be between 0 and 3. Actual: %d", index)
	}
	offset := KeyAreaOffset + 16*index
	return h.HeaderData[offset : offset+16], nil
}

func (h *NcaHeader) GetDecryptKey(index int) ([]byte, error) {
	return h.GetEncryptedKey(index)
}

func (h *NcaHeader) GetFsHeader(index int32) (*NcaFsHeader, error) {
	expectedHash := h.GetFsHeaderHash(index)
	offset := FsHeadersOffset + FsHeaderSize*index
	headerData := h.HeaderData[offset : offset+FsHeaderSize]
	actualHash := crypto.Sha256(headerData)
	if !tools.ArrayEquals(expectedHash, actualHash) {
		return nil, errors.WithStack(InvalidFsHeaderHashErr)
	}
	return NewNcaFsHeader(headerData), nil
}

func (h *NcaHeader) GetKeyGeneration() byte {
	if h.KeyGeneration1 > h.KeyGeneration2 {
		return h.KeyGeneration1
	} else {
		return h.KeyGeneration2
	}
}

func (h *NcaHeader) GetKeyArea() []byte {
	return h.HeaderData[KeyAreaOffset : KeyAreaOffset+KeyAreaSize]
}

func (h *NcaHeader) IsNca0() bool {
	return h.FormatVersion >= NcaVersionNca0
}

func (h *NcaHeader) BlockToOffset(blockIndex int32) int64 {
	return int64(blockIndex) * 0x200
}

func (h *NcaHeader) GetSectionSize(index int32) int64 {
	info := h.GetSectionEntry(index)
	return h.BlockToOffset(info.EndBlock - info.StartBlock)
}

func decryptNcaHeader(keySet *KeySet, reader fssystem.IStorage) ([]byte, bool, error) {
	buf, err := reader.ReadBytes(NcaHeaderSize)
	if err != nil {
		return nil, false, err
	}

	if CheckNcaHeaderIfDecrypted(buf) {
		decVersion := buf[0x203] - '0'

		if decVersion != 0 && decVersion != 2 && decVersion != 3 {
			return nil, false, errors.Errorf("NCA version %d is not supported", decVersion)
		}
		return buf, false, nil
	}

	key1 := keySet.HeaderKey.SubKeys(0)
	key2 := keySet.HeaderKey.SubKeys(1)

	transform, err := NewAes128XtsTransform(key1, key2, true)
	if err != nil {
		return nil, false, err
	}
	transform.TransformBlock(buf[:NcaHeaderSectorSize], 0)
	offset := NcaHeaderSectorSize * 1
	transform.TransformBlock(buf[offset:offset+NcaHeaderSectorSize], 1)

	if buf[0x200] != 'N' || buf[0x201] != 'C' || buf[0x202] != 'A' {
		return nil, false, errors.WithStack(NcaHeaderKeyIncorrectErr)
	}

	version := buf[0x203] - '0'

	if version == 3 {
		for sector := 2; sector < NcaHeaderSize/NcaHeaderSectorSize; sector++ {
			offset = sector * NcaHeaderSectorSize
			transform.TransformBlock(buf[offset:offset+NcaHeaderSectorSize], uint64(sector))
		}
	} else if version == 2 {
		for i := 0x400; i < NcaHeaderSize; i += NcaHeaderSectorSize {
			transform.TransformBlock(buf[i:i+NcaHeaderSectorSize], 0)
		}
	} else if version != 0 {
		return nil, false, errors.Errorf("NCA version %d is not supported", version)
	}

	return buf, true, nil
}

func CheckNcaHeaderIfDecrypted(header []byte) bool {
	// Check the magic value
	if header[0x200] != 'N' || header[0x201] != 'C' || header[0x202] != 'A' {
		return false
	}

	// Check the version in the magic value
	if header[0x203] >= '0' && header[0x203] <= '9' {
		return false
	}

	// Is the distribution type valid?
	if header[0x204] > byte(DistributionTypeGameCard) {
		return false
	}

	// Is the content type valid?
	if header[0x205] > byte(NcaContentTypePublicData) {
		return false
	}

	return true
}

func detectNcaVersion(header []byte) NcaVersion {
	version := header[0x203] - '0'

	if version == 3 {
		return NcaVersionNca3
	}
	if version == 2 {
		return NcaVersionNca2
	}
	if version != 0 {
		return NcaVersionUnknown
	}

	// There are multiple versions of NCA0 that each encrypt the key area differently.
	// Examine the key area to detect which version this NCA is.
	keyArea := header[KeyAreaOffset : KeyAreaOffset+KeyAreaSize]

	// The end of the key area will only be non-zero if it's RSA-OAEP encrypted
	if !tools.IsZeroArray(keyArea[0x80 : 0x80+0x10]) {
		return NcaVersionNca0RsaOaep
	}

	// Key areas using fixed, unencrypted keys always use the same keys.
	// Check for these keys by comparing the key area with the known hash of the fixed body keys.
	hash := crypto.Sha256(keyArea[:0x20])
	if tools.ArrayEquals(Nca0FixedBodyKeySha256Hash, hash) {
		return NcaVersionNca0FixedKey
	}

	// Otherwise the key area is encrypted the same as modern NCAs.
	return NcaVersionNca0
}

func NewNcaHeader(keySet *KeySet, reader fssystem.IStorage) (*NcaHeader, error) {
	header, encrypted, err := decryptNcaHeader(keySet, reader)
	if err != nil {
		return nil, err
	}
	ncaHeader := &NcaHeader{
		IsEncrypted:      encrypted,
		Signature1:       header[:0x100],
		Signature2:       header[0x100:0x200],
		Magic:            string(header[0x200 : 0x200+4]),
		DistributionType: DistributionType(header[0x204]),
		ContentType:      NcaContentType(header[0x205]),
		KeyGeneration1:   header[0x206],
		KeyAreaKeyIndex:  header[0x207],
		NcaSize:          binary.LittleEndian.Uint64(header[0x208:0x210]),
		TitleId:          binary.LittleEndian.Uint64(header[0x210:0x218]),
		SdkVersion:       binary.LittleEndian.Uint32(header[0x218:0x220]),
		KeyGeneration2:   header[0x220],
		RightsId:         header[0x230 : 0x230+0x10],
		FormatVersion:    detectNcaVersion(header),
		HeaderData:       header,
	}
	return ncaHeader, nil
}
